import os
import json
from py2neo import Graph, Node, Relationship
from tqdm import tqdm

class MedicalGraph:
    def __init__(self, db_name="medical"):
        """初始化医疗知识图谱，设置数据路径和Neo4j连接"""
        cur_dir = os.path.dirname(os.path.abspath(__file__))
        self.data_path = os.path.join(cur_dir, 'data/medical.json')
        self.db_name = db_name
        try:
            # 连接到 Neo4j，默认使用 system 数据库检查和创建新数据库
            self.g = Graph("bolt://localhost:7687", auth=("neo4j", "CHANG123"))
            # 检查数据库是否存在
            result = self.g.run("SHOW DATABASES").data()
            db_exists = any(db['name'] == self.db_name for db in result)
            if not db_exists:
                print(f"数据库 {self.db_name} 不存在，正在创建...")
                self.g.run(f"CREATE DATABASE {self.db_name}")
                print(f"数据库 {self.db_name} 创建成功")
            # 切换到指定数据库
            self.g = Graph("bolt://localhost:7687", auth=("neo4j", "CHANG123"), name=self.db_name)
        except Exception as e:
            raise ConnectionError(f"无法连接到Neo4j数据库或创建数据库失败: {e}")

    def read_nodes(self):
        """读取JSON文件，解析节点和关系数据"""
        # 初始化节点和关系列表
        nodes = {
            'drugs': set(), 'foods': set(), 'checks': set(),
            'departments': set(), 'producers': set(),
            'symptoms': set(), 'diseases': set()
        }
        disease_infos = []
        rels = {
            'department': [], 'noteat': [], 'doeat': [], 'recommandeat': [],
            'commondrug': [], 'recommanddrug': [], 'check': [],
            'drug_producer': [], 'symptom': [], 'acompany': [], 'category': []
        }

        with open(self.data_path, encoding='utf-8') as file:
            for line in file:
                data = json.loads(line)
                disease = {'name': data['name'], 'desc': '', 'prevent': '', 'cause': '',
                          'easy_get': '', 'cure_department': '', 'cure_way': '',
                          'cure_lasttime': '', 'symptom': '', 'cured_prob': ''}

                # 解析疾病相关信息
                for key in ['desc', 'prevent', 'cause', 'easy_get', 'cure_way', 'cure_lasttime', 'cured_prob']:
                    disease[key] = data.get(key, '')

                if 'get_prob' in data:
                    disease['easy_get'] = data['get_prob']

                # 处理症状
                if 'symptom' in data:
                    nodes['symptoms'].update(data['symptom'])
                    rels['symptom'].extend([[disease['name'], s] for s in data['symptom']])

                # 处理并发症
                if 'acompany' in data:
                    rels['acompany'].extend([[disease['name'], a] for a in data['acompany']])

                # 处理科室
                if 'cure_department' in data:
                    cure_department = data['cure_department']
                    disease['cure_department'] = cure_department
                    nodes['departments'].update(cure_department)
                    if len(cure_department) == 1:
                        rels['category'].append([disease['name'], cure_department[0]])
                    elif len(cure_department) == 2:
                        rels['department'].append([cure_department[1], cure_department[0]])
                        rels['category'].append([disease['name'], cure_department[1]])

                # 处理药物
                for drug_type in ['common_drug', 'recommand_drug']:
                    if drug_type in data:
                        nodes['drugs'].update(data[drug_type])
                        rels[drug_type.replace('_', '')].extend([[disease['name'], d] for d in data[drug_type]])

                # 处理食物
                for food_type in ['not_eat', 'do_eat', 'recommand_eat']:
                    if food_type in data:
                        nodes['foods'].update(data[food_type])
                        rels[food_type.replace('_', '')].extend([[disease['name'], f] for f in data[food_type]])

                # 处理检查
                if 'check' in data:
                    nodes['checks'].update(data['check'])
                    rels['check'].extend([[disease['name'], c] for c in data['check']])

                # 处理生产商
                if 'drug_detail' in data:
                    for detail in data['drug_detail']:
                        producer, drug = detail.split('(')
                        drug = drug.rstrip(')')
                        nodes['producers'].add(producer)
                        nodes['drugs'].add(drug)
                        rels['drug_producer'].append([producer, drug])

                nodes['diseases'].add(disease['name'])
                disease_infos.append(disease)

        return nodes, disease_infos, rels

    def create_node(self, label, nodes):
        """创建节点，检查是否已存在"""
        for node_name in tqdm(nodes, desc=f"创建 {label} 节点"):
            if not self.g.nodes.match(label, name=node_name).exists():
                node = Node(label, name=node_name)
                self.g.create(node)

    def create_diseases_nodes(self, disease_infos):
        """创建疾病节点，包含详细信息"""
        for info in tqdm(disease_infos, desc="创建 Disease 节点"):
            if not self.g.nodes.match("Disease", name=info['name']).exists():
                node = Node("Disease", **info)
                self.g.create(node)

    def create_relationship(self, start_label, end_label, edges, rel_type, rel_name):
        """创建关系，检查是否已存在"""
        for edge in tqdm(set(tuple(e) for e in edges), desc=f"创建 {rel_type} 关系"):
            start, end = edge
            query = (
                f"MATCH (p:{start_label} {{name: $start}}), (q:{end_label} {{name: $end}}) "
                f"MERGE (p)-[r:{rel_type} {{name: $rel_name}}]->(q)"
            )
            try:
                self.g.run(query, start=start, end=end, rel_name=rel_name)
            except Exception as e:
                print(f"创建 {rel_type} 关系 {start} -> {end} 失败: {e}")

    def create_graphnodes(self):
        """创建所有节点"""
        nodes, disease_infos, _ = self.read_nodes()
        print("开始创建知识图谱节点")
        self.create_diseases_nodes(disease_infos)
        for label, node_set in nodes.items():
            if label != 'diseases':
                self.create_node(label.capitalize(), node_set)
        print("完成知识图谱节点创建")

    def create_graphrels(self):
        """创建所有关系"""
        _, _, rels = self.read_nodes()
        print("开始创建知识图谱关系")
        relationships = [
            ('Disease', 'Food', rels['recommandeat'], 'recommand_eat', '推荐食谱'),
            ('Disease', 'Food', rels['noteat'], 'no_eat', '忌吃'),
            ('Disease', 'Food', rels['doeat'], 'do_eat', '宜吃'),
            ('Department', 'Department', rels['department'], 'belongs_to', '属于'),
            ('Disease', 'Drug', rels['commondrug'], 'common_drug', '常用药品'),
            ('Producer', 'Drug', rels['drug_producer'], 'drugs_of', '生产药品'),
            ('Disease', 'Drug', rels['recommanddrug'], 'recommand_drug', '好评药品'),
            ('Disease', 'Check', rels['check'], 'need_check', '诊断检查'),
            ('Disease', 'Symptom', rels['symptom'], 'has_symptom', '症状'),
            ('Disease', 'Disease', rels['acompany'], 'acompany_with', '并发症'),
            ('Disease', 'Department', rels['category'], 'belongs_to', '所属科室')
        ]
        for start, end, edges, rel_type, rel_name in relationships:
            self.create_relationship(start, end, edges, rel_type, rel_name)
        print("完成知识图谱关系创建")

    def export_data(self):
        """导出节点数据到文件"""
        nodes, _, _ = self.read_nodes()
        for node_type, node_set in nodes.items():
            with open(f"{node_type}.txt", 'w', encoding='utf-8') as f:
                f.write('\n'.join(node_set))

if __name__ == '__main__':
    try:
        # 指定数据库名称
        handler = MedicalGraph(db_name="medical_knowledge")
        print("步骤 1：导入图谱节点")
        handler.create_graphnodes()
        print("步骤 2：导入图谱关系")
        handler.create_graphrels()
        print("知识图谱创建成功")
    except Exception as e:
        print(f"创建知识图谱失败: {e}")